import torch
import torchvision
import torch.nn as nn

# VGG architecture, used for the perceptual loss using a pretrained VGG network
class Vgg19(torch.nn.Module):
    def __init__(self, requires_grad=False):
        super().__init__()
        vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(2, 7):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(7, 12):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(12, 21):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(21, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h_relu1 = self.slice1(X)
        h_relu2 = self.slice2(h_relu1)
        h_relu3 = self.slice3(h_relu2)
        h_relu4 = self.slice4(h_relu3)
        h_relu5 = self.slice5(h_relu4)
        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
        return out


# Perceptual loss that uses a pretrained VGG network
class VggLoss(nn.Module):
    def __init__(self, gpu_ids):
        super(VggLoss, self).__init__()
        if len(gpu_ids) > 0:
            self.vgg = Vgg19().cuda()
        else:
            self.vgg = Vgg19()
        self.criterion = nn.L1Loss()
        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]

    def forward(self, x, y):
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
        return loss


class CalcContentReltLoss(nn.Module):
    """Calc Content Relt Loss.
    """

    def __init__(self):
        super(CalcContentReltLoss, self).__init__()

    def forward(self, pred, target):
        """Forward Function.

        Args:
            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
        """
        dM = 1.
        Mx = self.calc_emd_loss(pred, pred)
        Mx = Mx / Mx.sum(1, keepdim=True)
        My = self.calc_emd_loss(target, target)
        My = My / My.sum(1, keepdim=True)
        loss_content = torch.abs(
            dM * (Mx - My)).mean() * pred.shape[2] * pred.shape[3]
        return loss_content

    def calc_emd_loss(self, pred, target):
        """calculate emd loss.

        Args:
            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
        """
        b, _, h, w = pred.shape
        # print("1", b, h, w)
        pred = pred.view(b, -1, w * h)
        # print("2", pred.shape)
        pred_norm = torch.sqrt((pred**2).sum(1).view(b, -1, 1))
        # print("3", pred_norm.shape)
        pred = pred.permute(0, 2, 1)
        # print("4", pred.shape)
        target_t = target.view(b, -1, w * h)
        # print("4-1", target_t.shape)
        target_norm = torch.sqrt((target**2).sum(1).view(b, 1, -1))
        # print("4-2", target_norm.shape)
        similarity = torch.bmm(pred, target_t) / pred_norm / target_norm
        # print("5", similarity)
        dist = 1. - similarity
        # print("6", dist)
        return dist


### content and style loss with VGG19
def calc_mean_std(feat, eps=1e-5):
    # eps is a small value added to the variance to avoid divide-by-zero.
    size = feat.size()
    assert (len(size) == 4)
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std


def build_vgg():
    vgg = nn.Sequential(
        nn.Conv2d(3, 3, (1, 1)),
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(3, 64, (3, 3)),
        nn.ReLU(),  # relu1-1
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(64, 64, (3, 3)),
        nn.ReLU(),  # relu1-2
        nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(64, 128, (3, 3)),
        nn.ReLU(),  # relu2-1
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(128, 128, (3, 3)),
        nn.ReLU(),  # relu2-2
        nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(128, 256, (3, 3)),
        nn.ReLU(),  # relu3-1
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(256, 256, (3, 3)),
        nn.ReLU(),  # relu3-2
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(256, 256, (3, 3)),
        nn.ReLU(),  # relu3-3
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(256, 256, (3, 3)),
        nn.ReLU(),  # relu3-4
        nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(256, 512, (3, 3)),
        nn.ReLU(),  # relu4-1
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu4-2
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu4-3
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu4-4
        nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu5-1, this is the last layer used
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu5-2
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU(),  # relu5-3
        nn.ReflectionPad2d((1, 1, 1, 1)),
        nn.Conv2d(512, 512, (3, 3)),
        nn.ReLU()  # relu5-4
    )
    return vgg


class VGG19(nn.Module):
    def __init__(self, checkpoint):
        super(VGG19, self).__init__()
        vgg = build_vgg()
        vgg.load_state_dict(torch.load(checkpoint))
        vgg_layers = list(vgg.children())
        self.enc_1 = nn.Sequential(*vgg_layers[:4])  # input -> relu1_1
        self.enc_2 = nn.Sequential(*vgg_layers[4:11])  # relu1_1 -> relu2_1
        self.enc_3 = nn.Sequential(*vgg_layers[11:18])  # relu2_1 -> relu3_1
        self.enc_4 = nn.Sequential(*vgg_layers[18:31])  # relu3_1 -> relu4_1
        self.enc_5 = nn.Sequential(*vgg_layers[31:45])  # relu4_1 -> relu5_1

        # fix the encoder
        for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4', 'enc_5']:
            for param in getattr(self, name).parameters():
                param.requires_grad = False

        self.mse_loss = nn.MSELoss()

    # extract relu1_1 - relu(n_layer)_1 from input image
    def encode_with_intermediate(self, x, n_layer=4):
        results = [x]
        for i in range(n_layer):
            func = getattr(self, 'enc_{:d}'.format(i + 1))
            results.append(func(results[-1]))
        return results[1:]

    # extract relu(n_layer)_1 from input image
    def encode(self, x, n_layer=4):
        for i in range(n_layer):
            x = getattr(self, 'enc_{:d}'.format(i + 1))(x)
        return x

    def calc_content_loss(self, input, target):
        assert (input.size() == target.size())
        assert (target.requires_grad is False)
        return self.mse_loss(input, target)

    def calc_style_loss(self, input, target):
        # assert (input.size() == target.size())
        assert (target.requires_grad is False)
        input_mean, input_std = calc_mean_std(input)
        target_mean, target_std = calc_mean_std(target)
        return self.mse_loss(input_mean, target_mean) + self.mse_loss(input_std, target_std)

    def forward(self, real_A, fake_B, real_B, fake_A=None, n_layer=4, content_weight=0.5, style_weight=1):

        real_A_feats = self.encode_with_intermediate(real_A, n_layer)
        real_B_feats = self.encode_with_intermediate(real_B, n_layer)
        if fake_A is not None:
            fake_A_feats = self.encode_with_intermediate(fake_A, n_layer)
        fake_B_feats = self.encode_with_intermediate(fake_B, n_layer)

        # content loss 
        loss_c = 0
        if content_weight > 0:
            loss_c = self.calc_content_loss(fake_B_feats[-1], real_A_feats[-1])
            if fake_A is not None:
                loss_c += self.calc_content_loss(fake_A_feats[-1], real_B_feats[-1])   # relu4_1
            loss_c = loss_c * content_weight

        # style loss 
        loss_s = 0
        if style_weight > 0:
            for i in range(0, n_layer):
                loss_s += self.calc_style_loss(fake_B_feats[i], real_B_feats[i])
                if fake_A is not None:
                    loss_s += self.calc_style_loss(fake_A_feats[i], real_A_feats[i])
   
        return loss_c, loss_s
    

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

import cv2
import numpy as np

import random

import torchvision.transforms as transforms
from PIL import Image


''' Optical flow warping function '''
def warp(x, flo, padding_mode='border'):
    B, C, H, W = x.size()

    # Mesh grid
    xx = torch.arange(0, W).view(1, -1).repeat(H, 1)
    yy = torch.arange(0, H).view(-1, 1).repeat(1, W)
    xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
    yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
    grid = torch.cat((xx, yy), 1).float()
    if x.is_cuda:
        grid = grid.cuda()
    vgrid = grid - flo

    # Scale grid to [-1,1]
    vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :] / max(W - 1, 1) - 1.0
    vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :] / max(H - 1, 1) - 1.0
    vgrid = vgrid.permute(0, 2, 3, 1)
    output = F.grid_sample(x, vgrid, padding_mode=padding_mode, mode='nearest')
    return output


''' Regularization from paper: Consistent Video Style Transfer via Compound Regularization'''
class TemporalLoss(nn.Module):
    def __init__(self, data_sigma=True, data_w=True, noise_level=0.001,
                 motion_level=8, shift_level=10):

        super(TemporalLoss, self).__init__()
        self.MSE = torch.nn.MSELoss()

        self.data_sigma = data_sigma
        self.data_w = data_w
        self.noise_level = noise_level
        self.motion_level = motion_level
        self.shift_level = shift_level

    """ Flow should have most values in the range of [-1, 1]. 
        For example, values x = -1, y = -1 is the left-top pixel of input, 
        and values  x = 1, y = 1 is the right-bottom pixel of input.
        Flow should be from pre_frame to cur_frame """

    def GaussianNoise(self, ins, mean=0, stddev=0.001):
        stddev = stddev + random.random() * stddev
        noise = Variable(ins.data.new(ins.size()).normal_(mean, stddev))

        if ins.is_cuda:
            noise = noise.cuda()
        return ins + noise

    def GenerateFakeFlow(self, height, width):
        ''' height: img.shape[0]
            width:  img.shape[1] '''

        if self.motion_level > 0:
            flow = np.random.normal(0, scale=self.motion_level, size=[height // 100, width // 100, 2])
            flow = cv2.resize(flow, (width, height))
            flow[:, :, 0] += random.randint(-self.shift_level, self.shift_level)
            flow[:, :, 1] += random.randint(-self.shift_level, self.shift_level)
            flow = cv2.blur(flow, (100, 100))
        else:
            flow = np.ones([width, height, 2])
            flow[:, :, 0] = random.randint(-self.shift_level, self.shift_level)
            flow[:, :, 1] = random.randint(-self.shift_level, self.shift_level)

        return torch.from_numpy(flow.transpose((2, 0, 1))).float()

    def GenerateFakeData(self, first_frame):
        ''' Input should be a (H,W,3) numpy of value range [0,1]. '''

        if self.data_w:
            forward_flow = self.GenerateFakeFlow(first_frame.shape[2], first_frame.shape[3])
            if first_frame.is_cuda:
                forward_flow = forward_flow.cuda()
            forward_flow = forward_flow.expand(first_frame.shape[0], 2, first_frame.shape[2], first_frame.shape[3])
            second_frame = warp(first_frame, forward_flow)
        else:
            second_frame = first_frame.clone()
            forward_flow = None

        if self.data_sigma:
            second_frame = self.GaussianNoise(second_frame, stddev=self.noise_level)

        return second_frame, forward_flow

    def forward(self, first_frame, second_frame, forward_flow):
        if self.data_w:
            first_frame = warp(first_frame, forward_flow)

        temporalloss = torch.mean(torch.abs(first_frame - second_frame))

        return temporalloss, first_frame
    

    def GenerateFakeDataByFlow(self, first_frame, flow):
        ''' Input should be a (1,3,H,W) tensor of value range [0,1]. '''

        forward_flow = flow
        if first_frame.is_cuda:
            forward_flow = forward_flow.cuda()
        print(first_frame.device, forward_flow.device)
        forward_flow = forward_flow.expand(first_frame.shape[0], 2, first_frame.shape[2], first_frame.shape[3])
        second_frame = warp(first_frame, forward_flow)

        second_frame = self.GaussianNoise(second_frame, stddev=self.noise_level)

        return second_frame, forward_flow